# -*- coding: utf-8 -*-
# time: 2025/4/30 16:05
# file: 多功能助手_Agent.py
# author: hanson
from langchain_community.utilities import WikipediaAPIWrapper
from datetime import datetime
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import ChatOllama
from langchain_community.tools import DuckDuckGoSearchRun
from langchain.tools import Tool, StructuredTool

# 1. 定义多个工具
# 1. 定义数学计算工具
def math_calculator(expression: str) -> str:
    """计算数学表达式，如 '2+2' 或 'sqrt(9)'"""
    try:
        result = eval(expression)*eval(expression) * 3.14  # 注意：实际生产环境应使用更安全的计算方式
        return str(result)
    except Exception as e:
        return f"计算错误: {e}"

# 数学工具（同上）
math_tool = Tool.from_function(
    name="math_calculator",
    description="计算数学表达式",
    func=math_calculator
)

# 维基百科工具
wikipedia = WikipediaAPIWrapper()
wikipedia_tool = Tool.from_function(
    name="wikipedia",
    description="查询维基百科获取事实信息",
    func=wikipedia.run
)

# 网络搜索工具
search = DuckDuckGoSearchRun()
search_tool = Tool.from_function(
    name="search",
    description="使用互联网搜索最新信息",
    func=search.run
)


# 时间工具
def get_current_time() -> str:
    """返回当前日期和时间"""
    try:
        return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    except Exception as e:
        return f"计算错误: {e}"
# 修正后的时间工具 - 使用StructuredTool明确不接受参数
time_tool =  StructuredTool.from_function(
    name="current_time",
    description="获取当前日期和时间",
    func=get_current_time
)

# 2. 创建 LLM
llm = ChatOllama(model="qwen2.5:1.5b", temperature=0.3)  # 稍高的温度增加创造性

# 3. 定义提示模板
prompt = ChatPromptTemplate.from_messages([
    ("system", """你是一个多功能助手，可以:
    - 回答一般知识问题
    - 进行数学计算
    - 查询最新信息
    - 提供当前时间

    请根据问题选择合适的工具，如果没有合适工具，尝试用你的知识回答。"""),
    ("user", "{input}"),
("placeholder", "{agent_scratchpad}") # 代理工作区（必须包含）
])

# 4. 创建 Agent 和执行器
tools = [math_tool, wikipedia_tool, search_tool, time_tool]
agent = create_tool_calling_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# 5. 使用示例
queries = [
    "今天的日期是什么？",
    "计算圆的面积，半径为5",
    "马克·扎克伯格的最新消息",
    "Python编程语言是谁创建的？"
]

print(datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
for query in queries:
    print(f"问题: {query}")
    result = agent_executor.invoke({"input": query})
    print(f"回答: {result['output']}\n")